from typing import Dict, List, Literal, Tuple, Union
import networkx as nx

from python_app.railsim.libapi import Track
from .tracks_model import TrackModel, TurnoutJunctionModel, Number


class Route:
    """
    一条进路
    """

    def __init__(
        self,
        name: str,
        direction: Literal["UP", "DOWN"],
        tracks: List[str],
        junction_states: Dict[str, Literal[0, 1]],
        final_action: Literal["STOP", "PASS"],
    ) -> None:
        self.direction = direction
        self.name = name
        self.final_action = final_action
        self.tracks = tracks
        self.junction_states = junction_states

    def get_track_objects(
        self, topology: nx.Graph
    ) -> List[Union[Track, Tuple[Number, Number], Tuple[Number, Number]]]:
        """
        获取轨道对象（C中用于和火车交互的对象）
        """
        track_objects = []
        for track_name in self.tracks:
            track: Union[TrackModel, TurnoutJunctionModel] = topology.nodes[track_name][
                "obj"
            ]
            if isinstance(track, TurnoutJunctionModel):
                state = self.junction_states[track_name]
                track_objects.append(track.get_track(state))
            else:
                track_objects.append(
                    (
                        track.track_object,
                        (track.start_x, track.start_y),
                        (track.end_x, track.end_y),
                    )
                )
        return track_objects

    def apply(self, topology: nx.Graph):
        for junction_name, state in self.junction_states.items():
            junction = topology.nodes[junction_name]["obj"]
            assert isinstance(junction, TurnoutJunctionModel)
            junction.state = state
        for i, track in enumerate(self.tracks):
            obj = topology.nodes[track]["obj"]
            if isinstance(obj, TrackModel):
                if self.direction == "UP":
                    obj.semaphore_up_state = "PASS"
                    obj.semaphore_down_state = "STOP"
                elif self.direction == "DOWN":
                    obj.semaphore_up_state = "STOP"
                    obj.semaphore_down_state = "PASS"
                else:
                    raise ValueError(f"Unknown direction: {self.direction}")
            elif isinstance(obj, TurnoutJunctionModel):
                pass
            else:
                raise ValueError(f"Unknown track type: {type(obj)}")
        last_track_obj = topology.nodes[self.tracks[-1]]["obj"]
        # 如果是最后一段轨道且要停车，就根据上行/下行开放信号
        if self.final_action == "STOP":
            if self.direction == "UP":
                last_track_obj.semaphore_up_state = "STOP"
            elif self.direction == "DOWN":
                last_track_obj.semaphore_down_state = "STOP"
        elif self.final_action == "PASS":
            if self.direction == "UP":
                last_track_obj.semaphore_up_state = "PASS"
            elif self.direction == "DOWN":
                last_track_obj.semaphore_down_state = "PASS"


class TrackTopology:
    def __init__(self) -> None:
        # track_model.draw(extra_surface)
        track1 = TrackModel((90, 100), (200, 100))
        junc1 = TurnoutJunctionModel((200, 100), (300, 100), (260, 130), 0)
        track2 = TrackModel((300, 100), (400, 100))
        track3 = TrackModel((260, 130), (400, 130))
        junc2 = TurnoutJunctionModel((400, 100), (500, 100), (400, 130), 1)
        track4 = TrackModel((500, 100), (600, 100))

        self.items = [track1, junc1, track3, track4, junc2, track2]
        self.topology = nx.DiGraph()
        self.topology.add_node("track#1", obj=track1)
        self.topology.add_node("track#main", obj=track2)
        self.topology.add_node("track#side", obj=track3)
        self.topology.add_node("track#4", obj=track4)
        self.topology.add_node("junc#1", obj=junc1)
        self.topology.add_node("junc#2", obj=junc2)

        self.topology.add_edge("track#1", "junc#1")
        self.topology.add_edge("junc#1", "track#main")
        self.topology.add_edge("junc#1", "track#side")
        self.topology.add_edge("track#main", "junc#2")
        self.topology.add_edge("track#side", "junc#2")
        self.topology.add_edge("junc#2", "track#4")

        nx.nx_pydot.write_dot(self.topology, "topology.dot")

        self.route1 = Route(
            "stop_sideline",
            "DOWN",
            ["track#1", "junc#1", "track#side"],
            dict([("junc#1", 1)]),
            "STOP",
        )
        self.route2 = Route(
            "pass_mainline",
            "DOWN",
            ["track#1", "junc#1", "track#main", "junc#2", "track#4"],
            dict([("junc#1", 0), ("junc#2", 0)]),
            "PASS",
        )
        self.route1.apply(self.topology)

        # self.track_objects: Dict[str, List[Track]] = {}
        # for node in self.topology.nodes:
        #     self.track_objects[node] = Track.create()

    def draw(self, surface):
        for node in self.topology.nodes:
            self.topology.nodes[node]["obj"].draw(surface)
